/*
 * stereocalib.cpp
 *
 *  Created on: Apr 15, 2013
 *      Author: rdu
 */

#include "stereocalib.h"

StereoCalib::StereoCalib()
{
	m_boardSize.width=8;
	m_boardSize.height=6;
	m_squareSize = 30;
	m_aspectRatio = 4.0/3.0;
	//cout<<"aspectRatio="<<m_aspectRatio<<endl;
	m_inputFilename = "stereo_image_list_rev.xml";
	//m_inputFilename = "stereo_image_list.xml";
	m_outputFilename = "out_camera_data.yml";
	m_displayCorners=false;
}

StereoCalib::~StereoCalib()
{

}

void StereoCalib::CalcChessboardCorners(Size boardSize, float squareSize, vector<Point3f>& corners)
{
	corners.resize(0);

	for( int i = 0; i < boardSize.height; i++ )
		for( int j = 0; j < boardSize.width; j++ )
			corners.push_back(Point3f(float(j*squareSize),
									  float(i*squareSize), 0));
}

bool StereoCalib::Run2Calibration(vector<vector<Point2f> > imagePoints1,
	                            vector<vector<Point2f> > imagePoints2,
	                            Size imageSize, Size boardSize,
	                            float squareSize, float aspectRatio,
	                            int flags,
	                            Mat& cameraMatrix1, Mat& distCoeffs1,
	                            Mat& cameraMatrix2, Mat& distCoeffs2,
	                            Mat& R12, Mat& T12)
{
	int c, i;

	// step 1: calibrate each camera individually
	vector<vector<Point3f> > objpt(1);
	vector<vector<Point2f> > imgpt;
	CalcChessboardCorners(boardSize, squareSize, objpt[0]);
	vector<Mat> rvecs, tvecs;

	for( c = 1; c <= 2; c++ )
	{
		const vector<vector<Point2f> >& imgpt0 = c == 1 ? imagePoints1 : imagePoints2;
		imgpt.clear();
		int N = 0;
		for( i = 0; i < (int)imgpt0.size(); i++ )
			if( !imgpt0[i].empty() )
			{
				imgpt.push_back(imgpt0[i]);
				N += (int)imgpt0[i].size();
			}

		if( imgpt.size() < 2 )
		{
			printf("Error: not enough views for camera %d\n", c);
			return false;
		}

		objpt.resize(imgpt.size(),objpt[0]);

		Mat cameraMatrix = Mat::eye(3, 3, CV_64F);
		if( flags & CV_CALIB_FIX_ASPECT_RATIO )
			cameraMatrix.at<double>(0,0) = aspectRatio;

		Mat distCoeffs = Mat::zeros(5, 1, CV_64F);

		double err = calibrateCamera(objpt, imgpt, imageSize, cameraMatrix,
						distCoeffs, rvecs, tvecs,
						flags|CV_CALIB_FIX_K3/*|CV_CALIB_FIX_K4|CV_CALIB_FIX_K5|CV_CALIB_FIX_K6*/);
		bool ok = checkRange(cameraMatrix) && checkRange(distCoeffs);
		if(!ok)
		{
			printf("Error: camera %d was not calibrated\n", c);
			return false;
		}
		printf("Camera %d calibration reprojection error = %g\n", c, sqrt(err/N));

		if( c == 1 )
			cameraMatrix1 = cameraMatrix, distCoeffs1 = distCoeffs;
		else
			cameraMatrix2 = cameraMatrix, distCoeffs2 = distCoeffs;
	}


	vector<vector<Point2f> > imgpt_right;

	// step 2: calibrate (1,2) pairs
	imgpt.clear();
	imgpt_right.clear();
	int N = 0;

	for( i = 0; i < (int)std::min(imagePoints1.size(), imagePoints2.size()); i++ )
		if( !imagePoints1.empty() && !imagePoints2[i].empty() )
		{
			imgpt.push_back(imagePoints1[i]);
			imgpt_right.push_back(imagePoints2[i]);
			N += (int)imagePoints2[i].size();
		}

	if( imgpt.size() < 2 )
	{
		printf("Error: not enough shared views for cameras 0 and %d\n", c);
		return false;
	}

	objpt.resize(imgpt.size(),objpt[0]);

	Mat R, T, E, F;
	double err = stereoCalibrate(objpt, imgpt, imgpt_right, cameraMatrix1, distCoeffs1,
								 cameraMatrix2, distCoeffs2,
								 imageSize, R, T, E, F,
								 TermCriteria(TermCriteria::COUNT, 30, 0),
								 CV_CALIB_FIX_INTRINSIC);

	printf("Pair (1,2) calibration reprojection error = %g\n", sqrt(err/(N*2)));

	R12 = R; T12 = T;

	return true;
}

bool StereoCalib::ReadStringList(const string& filename, vector<string>& l)
{
	l.resize(0);
	FileStorage fs(filename, FileStorage::READ);
	if( !fs.isOpened() )
		return false;
	FileNode n = fs.getFirstTopLevelNode();
	if( n.type() != FileNode::SEQ )
		return false;
	FileNodeIterator it = n.begin(), it_end = n.end();
	for( ; it != it_end; ++it )
		l.push_back((string)*it);
	return true;
}

int StereoCalib::StereoCalibration()
{
	int i, k;
	int flags = 0;
	Size imageSize;

	vector<vector<Point2f> > imgpt[2];
	vector<string> imageList;

	//flags |= CV_CALIB_FIX_ASPECT_RATIO;
	//flags |= CV_CALIB_ZERO_TANGENT_DIST;
	//flags |= CV_CALIB_FIX_PRINCIPAL_POINT;

	if( m_inputFilename.empty() ||
	   !ReadStringList(m_inputFilename, imageList) ||
	   imageList.size() == 0 || imageList.size() % 2 != 0 )
	{
		cout<<"imageList.size="<<imageList.size()<<endl;
		//printf("Error: the input image list is not specified, or can not be read, or the number of files is not divisible by 2\n");
		return -1;
	}

//	cout<<"size="<<imageList.size()<<endl;
//	cout<<"1="<<imageList[0]<<endl;
//	cout<<"2="<<imageList[1]<<endl;
//	cout<<"3="<<imageList[2]<<endl;
//	cout<<"4="<<imageList[3]<<endl;

	//initialize parameters
	Mat view, viewGray;
	Mat cameraMatrix[2], distCoeffs[2],  R12, T12;//,R[2], P[2];
	for( k = 0; k < 2; k++ )
	{
		cameraMatrix[k] = Mat_<double>::eye(3,3);
		cameraMatrix[k].at<double>(0,0) = m_aspectRatio;
		cameraMatrix[k].at<double>(1,1) = 1;
		distCoeffs[k] = Mat_<double>::zeros(5,1);
	}
	Mat R13=Mat_<double>::eye(3,3), T13=Mat_<double>::zeros(3,1);

	FileStorage fs;
//    namedWindow( "Image View", 0 );

	//define the size of each image array vector
	for( k = 0; k < 2; k++ )
		imgpt[k].resize(imageList.size()/2);

	int j=0;
	//bool displayCorners = true;//true;
	const int maxScale = 2;

	for( i = j = 0; i < (int)(imageList.size()/2); i++ )
	{
		for( k = 0; k < 2; k++ )
		{
			const string& filename = imageList[i*2+k];
			Mat img = imread(filename, 0);
			cout<<"filename="<<filename<<endl;
			//imshow("img", img);
			if(img.empty())
				break;
			if( imageSize == Size() )      //store the size of the first image
				imageSize = img.size();
			else if( img.size() != imageSize )
			{
				cout << "The image " << filename << " has the size different from the first image size. Skipping the pair\n";
				break;
			}
			bool found = false;
			vector<Point2f>& corners = imgpt[k][j];
			for( int scale = 1; scale <= maxScale; scale++ )
			{
				Mat timg;
				if( scale == 1 )
					timg = img;
				else
					resize(img, timg, Size(), scale, scale);
				found = findChessboardCorners(timg, m_boardSize, corners,
					CV_CALIB_CB_ADAPTIVE_THRESH | CV_CALIB_CB_NORMALIZE_IMAGE);
				if( found )
				{
					if( scale > 1 )
					{
						Mat cornersMat(corners);
						cornersMat *= 1./scale;
					}
					break;
				}
			}
			if( m_displayCorners )
			{
				cout << filename << endl;
				Mat cimg, cimg1;
				cvtColor(img, cimg, CV_GRAY2BGR);
				drawChessboardCorners(cimg, m_boardSize, corners, found);
				double sf = 640./MAX(img.rows, img.cols);
				resize(cimg, cimg1, Size(), sf, sf);
				imshow("corners", cimg1);
				char c = (char)waitKey(500);
				if( c == 27 || c == 'q' || c == 'Q' ) //Allow ESC to quit
					exit(-1);
			}
//			else
//				putchar('.');
			if( !found )
				break;
			cornerSubPix(img, corners, Size(11,11), Size(-1,-1),
						 TermCriteria(CV_TERMCRIT_ITER+CV_TERMCRIT_EPS,
									  30, 0.01));
		}

		if( k == 2 )j++;
	}

	cout << j << " pairs have been successfully detected.\n";

	if( j < 2 )
	{
		//cout << "Error: too little pairs to run the calibration\n";
		return -2;
	}

	printf("Running calibration ...\n");

	Run2Calibration(imgpt[0], imgpt[1], imageSize,
					m_boardSize, m_squareSize, m_aspectRatio, flags|CV_CALIB_FIX_K4|CV_CALIB_FIX_K5,
					cameraMatrix[0], distCoeffs[0],
					cameraMatrix[1], distCoeffs[1],
					R12, T12);

	fs.open(m_outputFilename, CV_STORAGE_WRITE);

	fs << "cameraMatrix1" << cameraMatrix[0];
	fs << "cameraMatrix2" << cameraMatrix[1];

	fs << "distCoeffs1" << distCoeffs[0];
	fs << "distCoeffs2" << distCoeffs[1];

	fs << "R12" << R12;
	fs << "T12" << T12;

	fs << "imageWidth" << imageSize.width;
	fs << "imageHeight" << imageSize.height;

	return 0;
}

